#include "KDTree.h"

#include <algorithm>
//#include <stdio>
//using namespace GM3;

unsigned int KDTree::thresh = 0;

KDTree::KDTree():
    G(NULL), root(NULL),
    treeSize(0), treeHeight(0), numLeaves(0), avgLeafSize(0.f), maxLeafSize(0)
{ }

KDTree::KDTree(GM3Graph *g):
    G(g), root(NULL),
    treeSize(0), treeHeight(0), numLeaves(0), avgLeafSize(0.f), maxLeafSize(0)
{ }

KDTree::~KDTree(){
    freeKDTree(root);
}

void KDTree::freeKDTree(KDTree::node *root){
    if(!root) return;

    treeSize = 0;
    treeHeight = 0;
    numLeaves = 0;
    avgLeafSize = 0.f;
    maxLeafSize = 0;

    freeKDTree(root->l);
    freeKDTree(root->r);
    delete root;
}

std::vector<unsigned int> KDTree::buildTree(){
    return buildTree(thresh);
}

std::vector<unsigned int> KDTree::buildTree(unsigned int thresh){
    using namespace std;
    if(!G) return vector<unsigned int>();
    vector<GM3Graph::node> nodesCopy(G->nodes);

    //misdirection vector to perform the rearranging operations on without
    //losing node IDs
    vector<unsigned int> misd(G->nodes.size());
    for(unsigned int i=0; i<misd.size(); i++) misd[i] = i;

    //clear out any old KD trees
    freeKDTree(root);

    //make a KD tree by calling the recursive function on the root
    root = buildSubTree(misd.begin(), misd.begin(), misd.end(), 0, thresh,
            boundingBox());

    //create a copy of the nodes vector to move them all in one step
    for(unsigned int i=0; i<G->nodes.size(); i++)
        G->nodes[i] = nodesCopy.at(misd[i]);

    //translation vector for rearranging the nodes and changing their IDs
    vector<unsigned int> trans(misd.size());
    for(unsigned int i=0; i<misd.size(); i++) trans.at(misd[i]) = i;


    //calculate average leaf size
    avgLeafSize = (float)G->numNodes() / (float)numLeaves;

    //translate the graph to the new node order
    G->translateEdges(trans);
    //G->translateParents(trans);

    return trans;
}


KDTree::node *KDTree::buildSubTree( std::vector<unsigned int>::iterator first,
        std::vector<unsigned int>::iterator s,
        std::vector<unsigned int>::iterator t,
        unsigned int d, unsigned int thresh, KDTree::boundingBox bb)
{
    using namespace std;
    unsigned int sz = (unsigned int)(t - s);

    if(sz <= 0){
        int  c = 10;
        return NULL;
    }

    node *n = new node(s-first, t-first);
    n->bb = bb;

    //statistics
    treeSize += 1;
    treeHeight = d+1 > treeHeight ? d+1 : treeHeight;


    if(sz <= thresh){ //this is a leaf node
        numLeaves++;
        maxLeafSize = sz > maxLeafSize ? sz : maxLeafSize;
        avgLeafSize += sz; //sum for now, and divide at the end
        n->p = findCenter(s, t);
        n->d = findRadius(s, t, n->p);
        n->w = sumWeight(s, t);
        return n;
    }

    float m = 0.5f;
    boundingBox bb0 = bb;
    boundingBox bb1 = bb;
    vector<unsigned int>::iterator s0 = s;
    vector<unsigned int>::iterator t0s1;
    vector<unsigned int>::iterator t1 = t;
    if(d%2 == 0){ //split on x
        if(d==6){
            int tmp = 17;
        }
        m = medianX(s, t);
//        if( m==0 ){
//            m = medianX(s, t);
//        }
        t0s1 = pivotSortX(s, t, m);
        bb0.ur.x(m);
        bb1.ll.x(m);
    }else{ //split on y
        m = medianY(s, t);
//        if( m==0 ){
//            m = medianY(s, t);
//        }
        t0s1 = pivotSortY(s, t, m);
        bb0.ur.y(m);
        bb1.ll.y(m);
    }

    //set up the child trees
    n->m = m;
//    if(t0s1 - s0 < 10 ||  t1 - t0s1 < 10){ // NICK DEBUG
//        std::cout << "zero size partition" << std::endl;
//    }
    n->l = buildSubTree(first, s0, t0s1, d+1, thresh, bb0);
    n->r = buildSubTree(first, t0s1, t1, d+1, thresh, bb1);

    bool childLnull = false;
    bool childRnull = false;
    if(n->l == NULL){
        cout << "[error] left child null" << endl;
        childLnull = true;
    }
    if(n->r == NULL){
        cout << "[error] right child null" << endl;
        childRnull = true;
    }


//    if(!childLnull && !childRnull){

        //determine center and radius from left and right child
        n->p = (n->l->p + n->r->p)/2.f;
        float dl = n->l->p.mag() + n->l->d;
        float dr = n->r->p.mag() + n->r->d;
        n->d = dr > dl ? dr : dl;

        //sum the weight of the two children
        n->w = n->l->w + n->r->w;

//    }
////    else{
//        if(childLnull && !childRnull){ //copy from right
//            //determine center and radius from left and right child
//            n->p = (n->r->p);// + n->r->p)/2.f;
//            //float dl = n->l->p.mag() + n->l->d;
//            float dr = n->r->p.mag() + n->r->d;
//            n->d = dr;// > dl ? dr : dl;

//            //sum the weight of the two children
//            n->w = n->r->w;// + n->r->w;

//        }
//        else if(  !childLnull && childRnull  ){ //copy from left
//            //determine center and radius from left and right child
//            n->p = (n->l->p);//
//            float dl = n->l->p.mag() + n->l->d;
//            //float dr = n->r->p.mag() + n->r->d;
//            n->d = dl;// > dl ? dr : dl;

//            //sum the weight of the two children
//            n->w = n->l->w;// + n->r->w;
//        }
//        else{
           // cout << "[error] both left and right are NULL!" << endl;
//        }
    //}

    return n;
}

cvec KDTree::findCenter(std::vector<unsigned int>::iterator s,
        std::vector<unsigned int>::iterator t){
    using namespace std;

    //just average all positions
    cvec avg = cvec();
    for(vector<unsigned int>::iterator i = s; i < t; i++){
        avg = avg + G->nodes[*i].p;
    }

    return avg / (float)(t-s);
}

float KDTree::findRadius(std::vector<unsigned int>::iterator s,
        std::vector<unsigned int>::iterator t, cvec c){
    using namespace std;

    float d = 0.f;
    float r = 0.f;
    for(vector<unsigned int>::iterator i = s; i < t; i++){
        d = (G->nodes[*i].p - c).mag();
        r = d > r ? d : r;
    }

    return r;
}

float KDTree::sumWeight(std::vector<unsigned int>::iterator s,
        std::vector<unsigned int>::iterator t){
    using namespace std;

    float w = 0.f;
    for(vector<unsigned int>::iterator i = s; i < t; i++){
        w += G->nodes[*i].weight;
    }

    return w;
}

std::vector<unsigned int>::iterator KDTree::pivotSortX(
        std::vector<unsigned int>::iterator s,
        std::vector<unsigned int>::iterator t, float m)
{
    unsigned int swp;
    std::vector<unsigned int>::iterator store = s;
    std::vector<unsigned int>::iterator i;
    for(i = s; i < t; i++){
        if(G->nodes[*i].p.x() < m){
            swp = *store;
            *store = *i;
            *i = swp;
            store++;
        }
    }

    return store;
//    //handle 0-size leaf
//    if( (store-s == 0) || (t-store == 0) ) return s+(t-s)/2; // 0415
//    else return store; // 0415
}

std::vector<unsigned int>::iterator KDTree::pivotSortY(
        std::vector<unsigned int>::iterator s,
        std::vector<unsigned int>::iterator t, float m)
{
    unsigned int swp;

    std::vector<unsigned int>::iterator store = s;
    std::vector<unsigned int>::iterator i;
    for(i = s; i < t; i++){
        if(G->nodes[*i].p.y() < m){
            swp = *store;
            *store = *i;
            *i = swp;
            store++;
        }
    }

    return store;

//    //handle 0-size leaf
//    if( (store-s == 0) || (t-store == 0) ) return s+(t-s)/2; // 0415
//    else return store; // 0415
}




bool compNodeX(GM3Graph::node n0, GM3Graph::node n1){
    float x0 = n0.p.x();
    float y0 = n1.p.x();
    return n0.p.x() < n1.p.x();
}
float KDTree::medianX(std::vector<unsigned int>::iterator s,
        std::vector<unsigned int>::iterator t)
{
    if(t-s > 1){
        float m;
        //get a set of nodes to choose a pivot from
        std::vector<GM3Graph::node> set = pivotSet(s, t);

        if(set.empty()){
            int tmp=17;
        }

        //sort the set of nodes
        sort(set.begin(), set.end(), compNodeX);

        unsigned int i = set.size()/2;
        m = (set[i].p.x() + set[i-1].p.x())/2.f;
        if(m != m){ //m is NaN
            std::cout << "X median calculation produced NaN at " << __FILE__ <<
                ":" << __LINE__ << std::endl;
        }
        return m;
    }else if(t-s > 0){
        int tmp = 17;
        return *s;
    }else{
        int tmp = 17;
        return 0.f;
    }
}


bool compNodeY(GM3Graph::node n0, GM3Graph::node n1){ return n0.p.y() < n1.p.y(); }
float KDTree::medianY(std::vector<unsigned int>::iterator s,
        std::vector<unsigned int>::iterator t)
{
    if(t-s > 1){
        float m;
        //get a set of nodes to choose a pivot from
        std::vector<GM3Graph::node> set = pivotSet(s, t);

        //sort the set of nodes
        sort(set.begin(), set.end(), compNodeY);

        unsigned int i = set.size()/2;
        m = (set[i].p.y() + set[i-1].p.y())/2.f; 
        if(m != m){ //m is NaN
            std::cout << "Y median calculation produced NaN at " << __FILE__ <<
                ":" << __LINE__ << std::endl;
        }
        return m;
    }else if(t-s > 0){
        return *s;
    }else{
        return 0.f;
    }
}


std::vector<GM3Graph::node> KDTree::pivotSet(
        std::vector<unsigned int>::iterator s,
        std::vector<unsigned int>::iterator t)
{
    using namespace std;
    unsigned int sz = (unsigned int)(t-s);
    vector<unsigned int> set;
    vector<GM3Graph::node> nodeSet;

    //if the desired space is bigger than the given number of nodes,
    //just use the whole set
    if(pivotSpaceSize >= sz){
        nodeSet.reserve(sz);
        for(vector<unsigned int>::iterator i=s; i<t; i++)
            nodeSet.push_back(G->nodes[*i]);
        return nodeSet;
    }

    //choose a set of unique random nodes to search for a pivot from
    set.reserve(pivotSpaceSize);
    nodeSet.reserve(pivotSpaceSize);
    srand(time(NULL));
    for(unsigned int i=0; i<pivotSpaceSize; i++){
        bool unique = false;
        unsigned int r;
        while(!unique){
            r = (unsigned int)rand() % sz;
            unique = true;
            vector<unsigned int>::iterator j = set.begin();
            while(j < set.end() && unique){ unique = !( *(j++) == r ); }
        }
        set.push_back( s[r] );
    }

    for(vector<unsigned int>::iterator i=set.begin(); i<set.end(); i++){
        nodeSet.push_back(G->nodes[*i]);
    }
    return nodeSet;
}


KDTree::node *KDTree::findLeaf(float x, float y){
    node *n = root;
    if(!n) return NULL;
    int d = 0;
    while(n->l && n->r){
        if( (d%2 == 0 && x <  n->m) || (d%2 == 1 && y <  n->m) ) n = n->l;
        else if( (d%2 == 0 && x >= n->m) || (d%2 == 1 && y >= n->m) ) n = n->r;
        d++;
    }

    return n;
}
